package net.cloudsun.tree.huffman;

import net.cloudsun.tree.binary.BinaryNode;
import net.cloudsun.tree.binary.BinaryTree;

import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.PriorityQueue;

/**
 * 赫夫曼树
 *
 * @author Ryan
 * @since 1.1
 */
public class HuffmanTree extends BinaryTree<Integer> implements java.lang.Comparable<HuffmanTree> {

    private Map<Character, Integer> frequency;
    private Map<Character, String> codeMap;
    private String result;
    public HuffmanTree(String text) {
        // 首先计算出各个字母出现的频率
        final char[] chars = text.toCharArray();
        this.frequency = new HashMap<>();
        for (char c : chars) {
            if (frequency.containsKey(c)) {
                frequency.put(c, frequency.get(c) + 1);
            } else {
                frequency.put(c, 1);
            }
        }
        // 组成一片森林，加入到优先队列里
        final PriorityQueue<HuffmanTree> forest = new PriorityQueue<>();
        frequency.forEach((key, value) -> forest.add(new HuffmanTree(key, value)));
        while (forest.size() > 1) {
           // 取出两个来然后加回去
            final HuffmanTree first = forest.poll();
            final HuffmanTree second = forest.poll();
            HuffmanTree merged = merge(first, second);
            forest.add(merged);
        }

        codeMap = new HashMap<>();
        this.root = forest.poll().root;

        this.bfs(x -> {
            HuffmanNode node = (HuffmanNode)x;
            if (((HuffmanNode) x).getElement() == null) {
                return true;
            }
            codeMap.put(node.getElement(), node.getFullCode());
            return true;
        });
        StringBuilder stringBuilder = new StringBuilder();
        for (char c : chars) {
            final String s = codeMap.get(c);
            stringBuilder.append(s);
        }
        this.result = stringBuilder.toString();
    }

    public HuffmanTree(int weight) {
        final HuffmanNode huffmanNode = new HuffmanNode(weight);
        setRoot(huffmanNode);
    }


    public HuffmanTree(Character key, Integer value) {
        final HuffmanNode huffmanNode = new HuffmanNode(value);
        huffmanNode.setElement(key);
        setRoot(huffmanNode);
    }

    @Override
    protected BinaryNode<Integer> createNode(Integer integer) {
        return new HuffmanNode(integer);
    }

    @Override
    public void delete(Integer integer) {
        throw new UnsupportedOperationException("不能删除节点");
    }

    @Override
    public int compareTo(HuffmanTree o) {
        if (o == null) {
            throw new NullPointerException();
        }
        return this.root.getValue().compareTo(o.root.getValue());
    }



    @Override
    public String toString() {
        return super.toString();
    }

    private static HuffmanTree merge(HuffmanTree first, HuffmanTree second) {
        final Integer weight1 = first.root.getValue();
        final Integer weight2 = second.root.getValue();
        int weight = weight1 + weight2;
        // 组成一个新树
        final HuffmanTree tree = new HuffmanTree(weight);
        tree.root.add(first.root);
        tree.root.add(second.root);
        return tree;
    }

    public Map<Character, Integer> getFrequency() {
        return frequency;
    }

    public Map<Character, String> getCodeMap() {
        return codeMap;
    }

    public String getResult() {
        return result;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (!(o instanceof HuffmanTree)) return false;
        HuffmanTree that = (HuffmanTree) o;
        return Objects.equals(this.root, that.root);
    }

    @Override
    public int hashCode() {
        return Objects.hash(root);
    }
}
